import * as tf from "@tensorflow/tfjs";
import * as tfd from "@tensorflow/tfjs-data";

class Webcam {
    webcam = null;
    truncatedMobileNet = null;
    model = null;

    // 初始化
    async init(video) {
        // 创建摄像头
        this.webcam = await tfd.webcam(video);

        // 加载模型
        const mobilenet = await tf.loadLayersModel("./webcam_model/model.json");
        const layer = mobilenet.getLayer("conv_pw_13_relu");
        this.truncatedMobileNet = tf.model({ inputs: mobilenet.inputs, outputs: layer.output });

        // 模型预热
        const screenShot = await this.webcam.capture();
        this.truncatedMobileNet.predict(screenShot.expandDims(0));
        screenShot.dispose();
    }

    // 获取实时图片
    async getImage() {
        const img = await this.webcam.capture();
        const processedImg = tf.tidy(() => img.expandDims(0).toFloat().div(127).sub(1));
        img.dispose();
        return processedImg;
    }

    // 是否收集图像中
    isCollecting = false;
    // 开始收集图像
    async startCollect(code, canvas) {
        if (this.isCollecting) return;
        this.isCollecting = true;
        while (this.isCollecting) {
            let img = await this.getImage();
            // 添加实例
            this.addExample(img, code);
            // 将图像绘制
            this.drawCanvas(img, canvas);
            // 销毁
            img.dispose();
            // 等待下一帧
            await tf.nextFrame();
        }
    }
    // 结束收集图像
    endCollect() {
        this.isCollecting = false;
    }
    // 清空收集图像
    clearCollect() {
        this.examples = null;
    }
    // 实例
    examples = null;
    // 标签
    labels = null;
    // 添加实例
    addExample(img, code) {
        // 获取实例和标签
        const example = this.truncatedMobileNet.predict(img);
        const label = tf.tidy(() => tf.oneHot(tf.tensor1d([code]).toInt(), 4));
        // 首次添加
        if (this.examples == null) {
            this.examples = tf.keep(example);
            this.labels = tf.keep(label);
        }
        // 非首次添加
        else {
            const oldExamples = this.examples;
            this.examples = tf.keep(oldExamples.concat(example, 0));

            const oldLabels = this.labels;
            this.labels = tf.keep(oldLabels.concat(label, 0));

            oldExamples.dispose();
            oldLabels.dispose();
            label.dispose();
        }
    }

    // 绘制图像到canvas展示
    drawCanvas(image, canvas) {
        const [width, height] = [224, 224];
        const ctx = canvas.getContext("2d");
        const imageData = new ImageData(width, height);
        const data = image.dataSync();
        for (let i = 0; i < height * width; ++i) {
            const j = i * 4;
            imageData.data[j + 0] = (data[i * 3 + 0] + 1) * 127;
            imageData.data[j + 1] = (data[i * 3 + 1] + 1) * 127;
            imageData.data[j + 2] = (data[i * 3 + 2] + 1) * 127;
            imageData.data[j + 3] = 255;
        }
        ctx.putImageData(imageData, 0, 0);
    }

    // 训练模型
    train() {
        if (this.examples == null) {
            alert("请添加实例");
            return;
        }
        this.model = tf.sequential({
            layers: [
                tf.layers.flatten({ inputShape: this.truncatedMobileNet.outputs[0].shape.slice(1) }),
                tf.layers.dense({
                    units: 100,
                    activation: "relu",
                    kernelInitializer: "varianceScaling",
                    useBias: true,
                }),
                tf.layers.dense({
                    units: 4,
                    activation: "softmax",
                    kernelInitializer: "varianceScaling",
                    useBias: false,
                }),
            ],
        });
        this.model.compile({ optimizer: tf.train.adam(0.0001), loss: "categoricalCrossentropy" });
        // 获取batchSize
        const batchSize = Math.floor(this.examples.shape[0] * 0.4);
        if (!(batchSize > 0)) {
            alert(`Batch size is 0 or NaN. Please choose a non-zero fraction.`);
            return;
        }
        // 开始训练
        return this.model.fit(this.examples, this.labels, {
            batchSize,
            epochs: 20,
        });
    }
    // 预测
    async predict(handle: (code: number) => void) {
        const img = await this.getImage();
        const embeddings = this.truncatedMobileNet.predict(img);
        const predictions = this.model.predict(embeddings);
        const predictedClass = predictions.as1D().argMax();
        const code = (await predictedClass.data())[0];
        img.dispose();
        handle(code);
        await tf.nextFrame();
    }
}

export default Webcam;
